import os
import json
import sys
import subprocess
from pathlib import Path

# 检查并安装第三方依赖
required_packages = ['requests', 'tqdm']
for package in required_packages:
    try:
        __import__(package)
    except ImportError:
        print(f"正在安装缺失依赖: {package}")
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])

import requests
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

def download_file(url, save_path):
    try:
        # 替换URL中的blob为resolve
        modified_url = url.replace("blob", "resolve")
        print(f"修改后的URL: {modified_url}")
        response = requests.get(modified_url, stream=True)
        response.raise_for_status()
        
        total_size = int(response.headers.get('content-length', 0))
        block_size = 1024
        
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        
        with open(save_path, 'wb') as f, tqdm(
            desc=os.path.basename(save_path),
            total=total_size,
            unit='B',
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for data in response.iter_content(block_size):
                f.write(data)
                bar.update(len(data))
        return True
    except Exception as e:
        print(f"下载失败 {url}: {str(e)}")
        return False

def main():
    # 加载模型列表
    script_dir = os.path.dirname(os.path.abspath(__file__))
    json_path = os.path.join(script_dir, 'reactormodellist.json')
    with open(json_path, 'r', encoding='utf-8') as f:
        model_list = json.load(f)
    
    # 准备下载任务
    tasks = []
    for folder, urls in model_list.items():
        for url in urls:
            filename = url.split('/')[-1]
            save_path = os.path.join(folder, filename)
            tasks.append((url, save_path))
    
    # 多线程下载
    with ThreadPoolExecutor(max_workers=5) as executor:
        results = list(executor.map(lambda args: download_file(*args), tasks))
    
    # 输出结果
    success = sum(results)
    print(f"\n下载完成: 成功 {success}/{len(results)} 个文件")

if __name__ == '__main__':
    main()